from io import BytesIO

import numpy as np
import torch
from PIL import Image


def convert_image_to_bytes(image):
    img = None
    for batch_number, image in enumerate(image):
        i = 255.0 * image.cpu().numpy()
        img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
    buffer = BytesIO()
    img.save(buffer, format="png")
    image_bytes = buffer.getvalue()
    return image_bytes


def convert_bytes_to_image(image_bytes):
    pil_image = Image.open(BytesIO(image_bytes))

    pil_image = pil_image.convert("RGB")
    pil_image = np.array(pil_image).astype(np.float32) / 255.0
    image = torch.from_numpy(pil_image)[None,]
    return image
